import os
import argparse
import numpy as np
import scipy.io as sio
from tqdm import tqdm
from scipy import signal
from scipy.interpolate import interp1d
import random

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--distortion_level", type=str, default="balanced", 
                       choices=["balanced"], help="失真程度 - 平衡级别，目标测试准确率50-60%")
    parser.add_argument("--noise_type", type=str, default="mixed", 
                       choices=["white", "pink", "brown", "mixed"], help="噪声类型")
    return parser.parse_args()

def generate_pink_noise(N, alpha=1.0):
    """生成粉色噪声 (1/f噪声)"""
    white_noise = np.random.randn(N)
    X = np.fft.fft(white_noise)
    
    # 创建1/f频谱
    freqs = np.fft.fftfreq(N)
    freqs[0] = 1e-10  # 避免除零
    
    # 应用1/f^alpha滤波
    S = 1 / (np.abs(freqs) ** (alpha/2))
    S[0] = S[1]  # 处理DC分量
    
    # 限制增益防止过大
    S = np.minimum(S, 100)
    
    # 应用滤波并转换回时域
    pink_noise = np.real(np.fft.ifft(X * S))
    return pink_noise / np.std(pink_noise)

def generate_brown_noise(N):
    """生成褐色噪声 (1/f²噪声)"""
    brown_noise = np.zeros(N)
    brown_noise[0] = np.random.randn()
    
    # 通过积分产生相关噪声
    for i in range(1, N):
        brown_noise[i] = 0.98 * brown_noise[i-1] + np.random.randn()
    
    return brown_noise / np.std(brown_noise)

def apply_frequency_domain_distortion(signal_data, distortion_level):
    """应用频域失真 - 极强版"""
    fft_data = np.fft.fft(signal_data)
    freq_mask = np.ones_like(fft_data)
    
    # 极大增加失真强度
    suppress_factor_range = (0.2, 0.5)   # 极强抑制
    enhance_factor_range = (2.0, 4.0)    # 极强增强
    noise_level = 0.25                   # 极高噪声水平
    
    # 抑制更多频率段
    for _ in range(4):  # 抑制4个不同频段
        suppress_start = int(len(fft_data) * np.random.uniform(0.05, 0.7))
        suppress_width = int(len(fft_data) * np.random.uniform(0.2, 0.4))
        suppress_end = min(suppress_start + suppress_width, len(fft_data))
        
        suppress_factor = np.random.uniform(*suppress_factor_range)
        freq_mask[suppress_start:suppress_end] = suppress_factor
        
        # 镜像到负频率
        if suppress_start > len(fft_data)//2:
            mirror_end = len(fft_data) - suppress_start + 1
            mirror_start = len(fft_data) - suppress_end + 1
            freq_mask[mirror_start:mirror_end] = suppress_factor
    
    # 增强更多频率段
    for _ in range(4):  # 增强4个不同频段
        enhance_start = int(len(fft_data) * np.random.uniform(0.05, 0.8))
        enhance_width = np.random.randint(40, 100)
        enhance_end = min(enhance_start + enhance_width, len(fft_data))
        
        enhance_factor = np.random.uniform(*enhance_factor_range)
        freq_mask[enhance_start:enhance_end] = enhance_factor
        
        # 镜像到负频率
        if enhance_start > len(fft_data)//2:
            mirror_end = len(fft_data) - enhance_start + 1
            mirror_start = len(fft_data) - enhance_end + 1
            freq_mask[mirror_start:mirror_end] = enhance_factor
    
    # 为全频段加噪声
    full_freq_noise = noise_level * np.std(np.abs(fft_data)) * \
                      (np.random.randn(len(fft_data)) + 
                       1j * np.random.randn(len(fft_data)))
    
    # 应用频域修改
    fft_data = fft_data * freq_mask + full_freq_noise
    
    # 确保频域对称性（实信号）
    if len(fft_data) % 2 == 0:
        fft_data[len(fft_data)//2+1:] = np.conj(np.flip(fft_data[1:len(fft_data)//2]))
    else:
        fft_data[(len(fft_data)+1)//2:] = np.conj(np.flip(fft_data[1:(len(fft_data)+1)//2]))
    
    return np.real(np.fft.ifft(fft_data))

def add_mechanical_distortions(signal_data, distortion_level):
    """添加机械失真 - 极强版"""
    t = np.arange(len(signal_data))
    distorted_signal = signal_data.copy()
    
    # 极大增加机械失真参数
    step_freq_range = (0.04, 0.15)
    step_amp_factor = 0.25  # 极强步进幅度
    drift_amp = 0.2         # 极强漂移幅度
    pattern_repeats = (4, 8)
    
    # 1. 极强步进信号
    step_freq = np.random.uniform(*step_freq_range)
    step_amplitude = step_amp_factor * np.std(signal_data)
    step_signal = step_amplitude * np.sign(np.sin(2 * np.pi * step_freq * t))
    # 添加更强的随机相位
    phase_noise = 0.5 * np.random.randn(len(t))
    step_signal *= (1 + 0.3 * np.sin(phase_noise))
    distorted_signal += step_signal
    
    # 2. 极复杂的频率漂移
    drift_factor = 1 + 0.02 * np.sin(2 * np.pi * t / len(t)) + \
                   0.015 * np.sin(4 * np.pi * t / len(t)) + \
                   0.01 * np.sin(6 * np.pi * t / len(t)) + \
                   0.008 * np.sin(8 * np.pi * t / len(t)) + \
                   0.005 * np.sin(10 * np.pi * t / len(t))
    
    t_original = np.arange(len(distorted_signal))
    t_drifted = np.cumsum(drift_factor)
    t_drifted = t_drifted * (len(t_original) / t_drifted[-1])
    
    # 使用插值实现频率漂移
    interp_func = interp1d(t_original, distorted_signal, kind='cubic', 
                          bounds_error=False, fill_value='extrapolate')
    distorted_signal = interp_func(t_drifted)
    
    # 3. 极多重复模式
    for _ in range(np.random.randint(3, 5)):
        pattern_length = np.random.randint(30, 100)
        pattern_start = np.random.randint(0, len(distorted_signal) - pattern_length)
        pattern = distorted_signal[pattern_start:pattern_start + pattern_length]
        
        num_repeats = np.random.randint(*pattern_repeats)
        for _ in range(num_repeats):
            insert_pos = np.random.randint(0, len(distorted_signal) - pattern_length)
            
            # 极强混合强度
            blend_factor = 0.9  # 90%混合
            distorted_signal[insert_pos:insert_pos + pattern_length] = \
                blend_factor * pattern + (1 - blend_factor) * \
                distorted_signal[insert_pos:insert_pos + pattern_length]
    
    return distorted_signal

def add_burst_noise(signal_data, distortion_level):
    """添加突发噪声"""
    # 大幅增加突发噪声
    num_bursts = min(6, max(3, len(signal_data) // 800))  # 更多突发
    burst_amp_factor = 2.5  # 更强的突发
    max_burst_length = 15   # 更长的突发
    
    distorted_signal = signal_data.copy()
    
    for _ in range(num_bursts):
        burst_loc = np.random.randint(0, len(signal_data) - max_burst_length)
        burst_length = np.random.randint(5, max_burst_length)
        
        # 生成多种形状的突发
        x = np.linspace(-2, 2, burst_length)
        burst_amp = burst_amp_factor * np.std(signal_data)
        
        # 随机选择突发形状
        burst_type = np.random.randint(0, 3)
        if burst_type == 0:
            burst_shape = burst_amp * np.exp(-x**2)  # 高斯
        elif burst_type == 1:
            burst_shape = burst_amp * np.sin(np.pi * x) * np.exp(-abs(x))  # 正弦包络
        else:
            burst_shape = burst_amp * np.random.uniform(-1, 1, burst_length)  # 随机
        
        distorted_signal[burst_loc:burst_loc + burst_length] += burst_shape
    
    return distorted_signal

def add_data_loss(signal_data, distortion_level):
    """添加数据丢失 - 极强版"""
    # 大幅增加数据丢失率
    loss_rate = 0.02  # 2%数据丢失
    
    missing_count = max(12, int(loss_rate * len(signal_data)))
    distorted_signal = signal_data.copy()
    
    # 更多更长的连续丢失
    for _ in range(missing_count//2):
        start_pos = np.random.randint(0, len(signal_data) - 12)
        length = np.random.randint(4, 12)
        end_pos = min(start_pos + length, len(signal_data))
        
        # 简单的线性插值加噪声
        if start_pos > 0 and end_pos < len(signal_data):
            lin_interp = np.linspace(distorted_signal[start_pos-1], 
                                   distorted_signal[end_pos], 
                                   end_pos - start_pos + 2)[1:-1]
            # 添加插值噪声
            interp_noise = 0.1 * np.std(distorted_signal) * np.random.randn(len(lin_interp))
            distorted_signal[start_pos:end_pos] = lin_interp + interp_noise
    
    return distorted_signal

def add_nonlinear_distortion(signal_data, distortion_level):
    """添加非线性失真 - 极强版"""
    # 极大增加非线性失真
    quad_factor = 0.08
    cubic_factor = 0.05
    fifth_factor = 0.02  # 添加五次项
    
    # 更严重的饱和效应
    saturation_threshold = 1.5 * np.std(signal_data)
    
    # 添加多阶失真项
    quad_term = quad_factor * signal_data ** 2
    cubic_term = cubic_factor * signal_data ** 3
    fifth_term = fifth_factor * signal_data ** 5
    
    distorted_signal = signal_data + quad_term - cubic_term + fifth_term
    
    # 应用双重软饱和
    distorted_signal = np.tanh(distorted_signal / saturation_threshold) * saturation_threshold
    distorted_signal = np.tanh(distorted_signal / (saturation_threshold * 0.8)) * saturation_threshold * 0.8
    
    return distorted_signal

def add_quantization(signal_data, distortion_level):
    """添加量化失真 - 极强版"""
    # 极大降低量化精度
    bits = np.random.randint(6, 9)  # 6-8位极低精度
    
    signal_range = np.max(signal_data) - np.min(signal_data)
    quant_levels = 2 ** bits
    quant_step = signal_range / quant_levels
    
    # 添加更强的量化噪声
    quant_noise = np.random.uniform(-quant_step, quant_step, len(signal_data))
    
    return np.round(signal_data / quant_step) * quant_step + quant_noise

def add_power_interference(signal_data, distortion_level):
    """添加电源干扰"""
    t = np.arange(len(signal_data))
    
    # 强烈的电源干扰
    power_freq = 50 + np.random.uniform(-5, 5)  # 45-55Hz
    interference_level = 0.08  # 强干扰
    
    # 添加多次谐波
    total_interference = 0
    for harmonic in [1, 2, 3]:
        harmonic_level = interference_level / harmonic
        harmonic_freq = power_freq * harmonic
        total_interference += harmonic_level * np.std(signal_data) * \
                            np.sin(2 * np.pi * harmonic_freq * t / 1000)
    
    return signal_data + total_interference

def apply_resonance_filter(signal_data, distortion_level):
    """应用共振滤波器"""
    # 强烈的共振滤波
    resonance_freq = np.random.uniform(0.1, 0.3)
    q_factor = np.random.uniform(8, 20)  # 高Q值
    
    # 创建IIR峰值滤波器
    nyquist = 0.5
    norm_freq = resonance_freq / nyquist
    
    # 使用butter滤波器近似共振滤波器
    b, a = signal.butter(4, [norm_freq * 0.8, norm_freq * 1.2], btype='band')
    
    # 应用滤波器
    filtered_signal = signal.filtfilt(b, a, signal_data)
    
    # 强烈混合原始信号和滤波信号
    mix_factor = 0.4  # 40%滤波信号
    return (1 - mix_factor) * signal_data + mix_factor * filtered_signal

def add_random_spikes(signal_data):
    """添加随机尖峰"""
    num_spikes = np.random.randint(5, 15)
    spike_amplitude = 3.0 * np.std(signal_data)
    
    distorted_signal = signal_data.copy()
    
    for _ in range(num_spikes):
        spike_pos = np.random.randint(0, len(signal_data))
        spike_width = np.random.randint(1, 4)
        
        for j in range(spike_width):
            if spike_pos + j < len(signal_data):
                distorted_signal[spike_pos + j] += spike_amplitude * np.random.uniform(-1, 1)
    
    return distorted_signal

def add_low_frequency_drift(signal_data):
    """添加低频漂移"""
    t = np.arange(len(signal_data))
    t_norm = t / len(t)
    
    # 复杂的低频漂移
    drift = 0.2 * np.std(signal_data) * (
        np.sin(2 * np.pi * 0.5 * t_norm) +
        0.5 * np.sin(2 * np.pi * 0.3 * t_norm) +
        0.3 * np.sin(2 * np.pi * 0.7 * t_norm)
    )
    
    return signal_data + drift

def generate_complex_noise(signal_data, noise_type, noise_level):
    """生成复杂噪声"""
    if noise_type == "white":
        noise = np.random.randn(len(signal_data))
    elif noise_type == "pink":
        noise = generate_pink_noise(len(signal_data))
    elif noise_type == "brown":
        noise = generate_brown_noise(len(signal_data))
    else:  # mixed
        # 混合不同类型的噪声
        white = np.random.randn(len(signal_data))
        pink = generate_pink_noise(len(signal_data))
        brown = generate_brown_noise(len(signal_data))
        
        # 随机权重混合
        weights = np.random.dirichlet([1, 1, 1])
        noise = weights[0] * white + weights[1] * pink + weights[2] * brown
    
    # 归一化并缩放
    noise = noise / np.std(noise) * noise_level * np.std(signal_data)
    return signal_data + noise

def simulate_comprehensive_distortion(signal_data, distortion_level, noise_type):
    """应用综合失真处理 - 极强版本"""
    # 确保信号是列向量
    if signal_data.ndim == 1:
        signal_data = signal_data.reshape(-1, 1).flatten()
    
    # 大幅增加噪声水平到20%
    noise_level = 0.20  # 20%噪声
    
    # 所有失真都必须应用（100%概率）
    apply_freq_dist = True    # 100%概率
    apply_mechanical = True   # 100%概率
    apply_quantization = True # 100%概率
    apply_resonance = True    # 100%概率
    apply_data_loss = True    # 100%概率
    apply_burst = True        # 100%概率
    apply_nonlinear = True    # 100%概率
    apply_power = True        # 100%概率
    apply_spikes = True       # 100%概率
    apply_lf_drift = True     # 100%概率
    
    distorted_signal = signal_data.copy()
    
    # 1. 频域失真
    distorted_signal = apply_frequency_domain_distortion(distorted_signal, distortion_level)
    
    # 2. 添加复杂噪声 (第一次)
    distorted_signal = generate_complex_noise(distorted_signal, noise_type, noise_level)
    
    # 3. 机械失真
    distorted_signal = add_mechanical_distortions(distorted_signal, distortion_level)
    
    # 4. 低频漂移
    distorted_signal = add_low_frequency_drift(distorted_signal)
    
    # 5. 量化失真
    distorted_signal = add_quantization(distorted_signal, distortion_level)
    
    # 6. 共振滤波
    distorted_signal = apply_resonance_filter(distorted_signal, distortion_level)
    
    # 7. 数据丢失
    distorted_signal = add_data_loss(distorted_signal, distortion_level)
    
    # 8. 突发噪声
    distorted_signal = add_burst_noise(distorted_signal, distortion_level)
    
    # 9. 随机尖峰
    distorted_signal = add_random_spikes(distorted_signal)
    
    # 10. 非线性失真
    distorted_signal = add_nonlinear_distortion(distorted_signal, distortion_level)
    
    # 11. 电源干扰
    distorted_signal = add_power_interference(distorted_signal, distortion_level)
    
    # 12. 添加第二轮复杂噪声
    extra_noise_level = 0.08
    distorted_signal = generate_complex_noise(distorted_signal, "mixed", extra_noise_level)
    
    # 13. 更强的最终缩放
    final_scale = np.random.uniform(0.6, 1.4)  # ±40%缩放
    distorted_signal *= final_scale
    
    # 14. 添加强化的周期性故障特征扰动
    t = np.arange(len(distorted_signal))
    periodic_distortion = 0.08 * np.std(distorted_signal) * (
        np.sin(2 * np.pi * t / 40) + 
        0.7 * np.sin(2 * np.pi * t / 60) +
        0.5 * np.sin(2 * np.pi * t / 80) +
        0.3 * np.sin(2 * np.pi * t / 100)
    )
    distorted_signal += periodic_distortion
    
    # 15. 添加高频毛刺
    spike_density = 0.02  # 2%的点有毛刺
    num_spikes = int(spike_density * len(distorted_signal))
    spike_positions = np.random.choice(len(distorted_signal), num_spikes, replace=False)
    for pos in spike_positions:
        distorted_signal[pos] += np.random.uniform(-2, 2) * np.std(distorted_signal) * 0.5
    
    # 16. 最终强随机扰动
    final_noise = 0.05 * np.std(distorted_signal) * np.random.randn(len(distorted_signal))
    distorted_signal += final_noise
    
    # 17. 添加阶跃失真
    if np.random.random() > 0.3:  # 70%概率
        step_position = np.random.randint(len(distorted_signal) // 4, 3 * len(distorted_signal) // 4)
        step_amplitude = np.random.uniform(0.5, 1.5) * np.std(distorted_signal)
        distorted_signal[step_position:] += step_amplitude * np.random.choice([-1, 1])
    
    return distorted_signal

def load_cwru_data(data_path):
    """只加载训练集数据用于生成仿真"""
    train_path = os.path.join(data_path, "train")
    if not os.path.exists(train_path):
        raise FileNotFoundError(f"训练数据目录 {train_path} 不存在，请先运行 preprocess.py")
    
    data_files = [f for f in os.listdir(train_path) if f.endswith('.npy')]
    processed_data = {}
    
    for file in tqdm(data_files, desc="加载训练数据文件"):
        fault_type = file.replace('.npy', '')
        file_path = os.path.join(train_path, file)
        segments = np.load(file_path)
        processed_data[fault_type] = segments
        print(f"加载 {fault_type} 训练数据，形状: {segments.shape}")
    
    return processed_data

def generate_simulated_data(original_data, args):
    """生成仿真数据"""
    simulated_data = {}
    
    for fault_type, segments in tqdm(original_data.items(), desc="生成仿真数据"):
        simulated_segments = []
        
        for segment in segments:
            # 应用综合失真
            simulated_segment = simulate_comprehensive_distortion(
                segment, args.distortion_level, args.noise_type)
            simulated_segments.append(simulated_segment)
        
        simulated_data[fault_type] = np.array(simulated_segments)
    
    return simulated_data

def save_processed_data(processed_data, save_path):
    """保存处理后的数据"""
    os.makedirs(save_path, exist_ok=True)
    
    for fault_type, segments in processed_data.items():
        # 将数据转换为numpy数组
        data_array = np.array(segments)
        
        # 保存为.npy格式
        save_file = os.path.join(save_path, f"{fault_type}.npy")
        np.save(save_file, data_array)
        print(f"已保存 {fault_type} 类别的数据，形状: {data_array.shape}")

def main():
    args = parse_args()
    
    # 设置路径 - 现在基于预处理后的训练集
    data_path = "./cwru_prepro"  # 包含train和test子目录的路径
    save_path = "./cwru_sim"  # 仿真数据保存路径
    
    # 只加载训练集数据
    print("加载CWRU训练集数据...")
    original_data = load_cwru_data(data_path)
    
    # 基于训练集生成仿真数据
    print(f"基于训练集生成仿真数据（失真程度：{args.distortion_level}，噪声类型：{args.noise_type}）...")
    simulated_data = generate_simulated_data(original_data, args)
    
    # 保存仿真数据
    print("保存仿真数据...")
    save_processed_data(simulated_data, save_path)
    print("仿真数据生成完成！")

if __name__ == "__main__":
    main() 